import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from model import ImageClassifier
from dataset import get_cifar10_data_loaders
from train import train_model, evaluate_model
from utils import save_model, load_model

def main():
    # 设置设备
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")

    # 获取数据加载器
    train_loader, test_loader = get_cifar10_data_loaders()

    # 初始化模型、损失函数和优化器
    model = ImageClassifier().to(device)
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=0.001)

    # 训练模型
    num_epochs = 50
    trained_model = train_model(model, train_loader, criterion, optimizer, device, num_epochs)

    # 评估模型
    accuracy = evaluate_model(trained_model, test_loader, device)
    print(f"Test Accuracy: {accuracy:.2f}%")

    # 保存模型
    save_model(trained_model, "image_classifier.pth")

if __name__ == "__main__":
    main()